import numpy as np
from const import *
import matplotlib.pyplot as plt

class NewTon1:
    def __init__(self, A, x0):
        self.A = A
        self.AT = np.transpose(A)
        self.x0 = x0
        self.n = A.shape[1]
        self.p = A.shape[0]

    def f(self, x):
        return x @ np.log(x)

    def gradient(self, x):
        return np.log(x) + 1

    def hessian(self, x):
        res = np.zeros((len(x), len(x)))
        for i in range(len(x)):
            res[i][i] = 1.0 / x[i]
        return res

    def dx(self, x):
        mat = np.vstack((np.hstack((self.hessian(x),self.AT)),np.hstack((self.A,np.zeros((self.p,self.p))))))
        right = np.hstack((-1 * self.gradient(x),np.zeros((self.p))))
        dy = np.linalg.inv(mat) @ right
        return dy[0:self.n]

    def declambda(self, x, dx):
        return dx @ self.hessian(x) @ dx

    def ok(self, x):
        for i in x:
            if i <= 0:
                return False
        return True

    def search(self):
        x0 = self.x0
        ls_f = [self.f(x0)]
        while True:
            t = 1
            dx = self.dx(x0)
            dec = self.declambda(x0, dx)
            if 0.5 * dec < epsilon:
                ls_f.append(self.f(x0))
                print(f"epoach {len(ls_f) - 1}  λ(x)^2: {dec}  f(x): {ls_f[-1]}")
                break
            x1 = x0 + t * dx
            while self.ok(x1) == False:
                t = t * beta
                x1 = x0 + t * dx
            while self.f(x1) > self.f(x0) - alpha * t * dec:
                t = beta * t
                x1 = x0 + t * dx
            x0 = x1
            ls_f.append(self.f(x0))
            print(f"epoach {len(ls_f) - 1}  λ(x)^2: {dec}  f(x): {ls_f[-1]}")
        return x0, ls_f

    def search_plot(self):
        x, ls = self.search()
        print(f"x={x}")
        print(f"f(x) = {ls[-1]}")
        plt.plot(range(len(ls)),ls)
        plt.xlabel("epoach")
        plt.ylabel("f(x)")
        plt.show()


if __name__ == "__main__":
    A = np.loadtxt("A.txt")
    x0 = np.loadtxt("x0.txt")
    new = NewTon1(A, x0)
    new.search_plot()
